# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

# -*- coding: utf-8 -*-
"""
测试全局函数的代码生成

验证 global_functions.py 中配置的全局函数能够正确生成 C++ 代码。
"""

from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains


# ============================================================================
# 数学函数测试
# ============================================================================


def test_abs_function():
    """测试 abs() 函数的代码生成"""

    def test_abs_int(x: int) -> int:
        """测试 abs 整数"""
        return abs(x)

    cpp = compile_chunk(test_abs_int)
    expected = """
    int test_abs_int(int x) {
        return std::abs(x);
    }
    """
    assert_code_contains(cpp, expected)

    def test_abs_float(x: float) -> float:
        """测试 abs 浮点数"""
        return abs(x)

    cpp = compile_chunk(test_abs_float)
    expected = """
    double test_abs_float(double x) {
        return std::abs(x);
    }
    """
    assert_code_contains(cpp, expected)


def test_pow_function():
    """测试 pow() 函数的代码生成"""

    def test_pow_int(x: int, y: int) -> int:
        """测试 pow 整数"""
        return pow(x, y)

    cpp = compile_chunk(test_pow_int)
    expected = """
    int test_pow_int(int x, int y) {
        return static_cast<int>(std::pow(x, y));
    }
    """
    assert_code_contains(cpp, expected)

    def test_pow_float(x: float, y: float) -> float:
        """测试 pow 浮点数"""
        return pow(x, y)

    cpp = compile_chunk(test_pow_float)
    expected = """
    double test_pow_float(double x, double y) {
        return std::pow(x, y);
    }
    """
    assert_code_contains(cpp, expected)


def test_min_variadic():
    """测试 min() 可变参数版本的代码生成"""

    def test_min_two(a: int, b: int) -> int:
        """测试 min 两个参数"""
        return min(a, b)

    cpp = compile_chunk(test_min_two)
    expected = """
    int test_min_two(int a, int b) {
        return mcpy::min(a, b);
    }
    """
    assert_code_contains(cpp, expected)

    def test_min_three(a: int, b: int, c: int) -> int:
        """测试 min 三个参数"""
        return min(a, b, c)

    cpp = compile_chunk(test_min_three)
    expected = """
    int test_min_three(int a, int b, int c) {
        return mcpy::min(a, b, c);
    }
    """
    assert_code_contains(cpp, expected)


def test_max_variadic():
    """测试 max() 可变参数版本的代码生成"""

    def test_max_two(a: int, b: int) -> int:
        """测试 max 两个参数"""
        return max(a, b)

    cpp = compile_chunk(test_max_two)
    expected = """
    int test_max_two(int a, int b) {
        return mcpy::max(a, b);
    }
    """
    assert_code_contains(cpp, expected)

    def test_max_three(a: float, b: float, c: float) -> float:
        """测试 max 三个参数"""
        return max(a, b, c)

    cpp = compile_chunk(test_max_three)
    expected = """
    double test_max_three(double a, double b, double c) {
        return mcpy::max(a, b, c);
    }
    """
    assert_code_contains(cpp, expected)


def test_min_list():
    """测试 min() 列表版本的代码生成（使用通用迭代器版本）"""

    def test_min_list(numbers: list[int]) -> int:
        """测试 min 列表"""
        return min(numbers)

    cpp = compile_chunk(test_min_list)
    expected = """
    int test_min_list(const mc::array<int>& numbers) {
        return mcpy::iterable_min<int>(numbers);
    }
    """
    assert_code_contains(cpp, expected)


def test_max_set():
    """测试 max() 集合版本的代码生成（使用通用迭代器版本）"""

    def test_max_set(numbers: set[int]) -> int:
        """测试 max 集合"""
        return max(numbers)

    cpp = compile_chunk(test_max_set)
    expected = """
    int test_max_set(const mcpy::variant_set& numbers) {
        return mcpy::iterable_max<int>(numbers);
    }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# 容器函数测试
# ============================================================================


def test_len_function():
    """测试 len() 函数的代码生成"""

    def test_len_list(items: list[int]) -> int:
        """测试 len 列表"""
        return len(items)

    cpp = compile_chunk(test_len_list)
    expected = """
    int test_len_list(const mc::array<int>& items) {
        return items.size();
    }
    """
    assert_code_contains(cpp, expected)

    def test_len_dict(data: dict[str, int]) -> int:
        """测试 len 字典"""
        return len(data)

    cpp = compile_chunk(test_len_dict)
    expected = """
    int test_len_dict(const mc::dict& data) {
        return data.size();
    }
    """
    assert_code_contains(cpp, expected)

    def test_len_set(items: set[str]) -> int:
        """测试 len 集合"""
        return len(items)

    cpp = compile_chunk(test_len_set)
    expected = """
    int test_len_set(const mcpy::variant_set& items) {
        return items.size();
    }
    """
    assert_code_contains(cpp, expected)


def test_repr_function():
    """测试 repr() 函数的代码生成"""

    def test_repr_list(items: list[int]) -> str:
        """测试 repr 列表"""
        return repr(items)

    cpp = compile_chunk(test_repr_list)
    expected = """
    std::string test_repr_list(const mc::array<int>& items) {
        return mcpy::list_repr(items);
    }
    """
    assert_code_contains(cpp, expected)

    def test_repr_dict(data: dict[str, int]) -> str:
        """测试 repr 字典"""
        return repr(data)

    cpp = compile_chunk(test_repr_dict)
    expected = """
    std::string test_repr_dict(const mc::dict& data) {
        return mcpy::dict_repr(data);
    }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# 序列函数测试
# ============================================================================


def test_enumerate_function():
    """测试 enumerate() 函数的代码生成"""

    code = """
def test_enumerate_list(items: list[str]) -> None:
    '''测试 enumerate 列表'''
    for idx, item in enumerate(items):
        pass
"""

    cpp = compile_chunk(code)
    expected = """
    void test_enumerate_list(const mc::array<std::string>& items) {
        for (const auto& [idx, item] : mcpy::enumerate(items)) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_enumerate_with_start():
    """测试 enumerate() 带 start 参数"""

    code = """
def test_enumerate_start(items: list[int]) -> None:
    '''测试 enumerate 带 start'''
    for idx, item in enumerate(items, start=10):
        pass
"""

    cpp = compile_chunk(code)
    expected = """
    void test_enumerate_start(const mc::array<int>& items) {
        for (const auto& [idx, item] : mcpy::enumerate(items, 10)) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_sorted_function():
    """测试 sorted() 函数的代码生成"""

    def test_sorted_list(items: list[int]) -> list[int]:
        """测试 sorted 列表"""
        return sorted(items)

    cpp = compile_chunk(test_sorted_list)
    expected = """
    mc::array<int> test_sorted_list(const mc::array<int>& items) {
        return mcpy::iterable_sorted<int>(items);
    }
    """
    assert_code_contains(cpp, expected)

    def test_sorted_reverse(items: list[int]) -> list[int]:
        """测试 sorted 反向"""
        return sorted(items, reverse=True)

    cpp = compile_chunk(test_sorted_reverse)
    expected = """
    mc::array<int> test_sorted_reverse(const mc::array<int>& items) {
        return mcpy::iterable_sorted<int>(items, true);
    }
    """
    assert_code_contains(cpp, expected)


def test_reversed_function():
    """测试 reversed() 函数的代码生成"""

    code = """
def test_reversed_list(items: list[int]) -> None:
    '''测试 reversed 列表'''
    for item in reversed(items):
        pass
"""

    cpp = compile_chunk(code)
    expected = """
    void test_reversed_list(const mc::array<int>& items) {
        for (const auto& item : mcpy::iterable_reversed(items)) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_sum_function():
    """测试 sum() 函数的代码生成"""

    def test_sum_list(numbers: list[int]) -> int:
        """测试 sum 列表"""
        return sum(numbers)

    cpp = compile_chunk(test_sum_list)
    # sum 现在使用通用的 reduce_sum（支持任何 iterable）
    assert "mcpy::reduce_sum" in cpp

    def test_sum_with_start(numbers: list[int]) -> int:
        """测试 sum 带初始值"""
        return sum(numbers, start=10)

    cpp = compile_chunk(test_sum_with_start)
    # sum 现在使用通用的 reduce_sum
    assert "mcpy::reduce_sum" in cpp


def test_any_all_functions():
    """测试 any() 和 all() 函数的代码生成"""

    def test_any_list(flags: list[bool]) -> bool:
        """测试 any 列表"""
        return any(flags)

    cpp = compile_chunk(test_any_list)
    expected = """
    bool test_any_list(const mc::array<bool>& flags) {
        return mcpy::iterable_any(flags);
    }
    """
    assert_code_contains(cpp, expected)

    def test_all_set(flags: set[bool]) -> bool:
        """测试 all 集合"""
        return all(flags)

    cpp = compile_chunk(test_all_set)
    # Protocol 版本统一使用 list_all（支持所有可迭代类型）
    expected = """
    bool test_all_set(const mcpy::variant_set& flags) {
        return mcpy::iterable_all(flags);
    }
    """
    assert_code_contains(cpp, expected)


def test_zip_function():
    """测试 zip() 函数的代码生成"""

    code = """
def test_zip_two_lists(a: list[int], b: list[str]) -> None:
    '''测试 zip 两个列表'''
    for x, y in zip(a, b):
        pass
"""

    cpp = compile_chunk(code)
    expected = """
    void test_zip_two_lists(const mc::array<int>& a, const mc::array<std::string>& b) {
        for (const auto& [x, y] : mcpy::zip_view(a, b)) {
        }
    }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# 其他函数测试
# ============================================================================


def test_print_function():
    """测试 print() 函数的代码生成"""

    def test_print_single(x: int) -> None:
        """测试 print 单个参数"""
        print(x)

    cpp = compile_chunk(test_print_single)
    expected = """
    void test_print_single(int x) {
        mcpy::print(x);
    }
    """
    assert_code_contains(cpp, expected)

    def test_print_multiple(a: int, b: str) -> None:
        """测试 print 多个参数"""
        print(a, b)

    cpp = compile_chunk(test_print_multiple)
    expected = """
    void test_print_multiple(int a, const std::string& b) {
        mcpy::print(a, b);
    }
    """
    assert_code_contains(cpp, expected)


def test_range_function():
    """测试 range() 函数在循环中的代码生成"""

    code = """
def test_range_loop() -> int:
    '''测试 range 循环'''
    total: int = 0
    for i in range(5):
        total = total + i
    return total
"""

    cpp = compile_chunk(code)
    expected = """
    int test_range_loop() {
        int total = 0;
        for (int i = 0; i < 5; ++i) {
            total = total + i;
        }
        return total;
    }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# 组合使用测试
# ============================================================================


def test_combined_global_functions():
    """测试全局函数的组合使用"""

    code = """
def process_numbers(numbers: list[int]) -> int:
    '''组合使用多个全局函数'''
    total: int = sum(numbers)
    count: int = len(numbers)
    if count > 0:
        avg: int = total // count
        return max(avg, min(numbers))
    return 0
"""

    cpp = compile_chunk(code)
    # 验证使用了通用的函数
    assert "mcpy::reduce_sum" in cpp  # sum 使用 reduce_sum
    assert "mcpy::iterable_min" in cpp  # min 使用 iterable_min
    assert "mcpy::max" in cpp


def test_nested_function_calls():
    """测试嵌套的全局函数调用"""

    def test_nested(items: list[int]) -> int:
        """测试嵌套调用"""
        return abs(min(items))

    cpp = compile_chunk(test_nested)
    expected = """
    int test_nested(const mc::array<int>& items) {
        return std::abs(mcpy::iterable_min<int>(items));
    }
    """
    assert_code_contains(cpp, expected)


def test_global_functions_in_loop():
    """测试在循环中使用全局函数"""

    code = """
def test_in_loop(data: list[list[int]]) -> list[int]:
    '''在循环中使用全局函数'''
    result: list[int] = []
    for sublist in data:
        if len(sublist) > 0:
            result.append(max(sublist))
    return result
"""

    cpp = compile_chunk(code)
    expected = """
    mc::array<int> test_in_loop(const mc::array<mc::array<int>>& data) {
        mc::array<int> result = {};
        for (const auto& sublist : data) {
            if (sublist.size() > 0) {
                result.push_back(mcpy::iterable_max<int>(sublist));
            }
        }
        return result;
    }
    """
    assert_code_contains(cpp, expected)
